from core.cuda import cuda_module
from core.function import Function
import utils

class SoftmaxCrossEntropy(Function):
    def forward(self, x, t):
        """
        前向传播函数，计算 Softmax Cross Entropy 损失。

        Parameters:
            x (Variable): 模型的输出结果。
            t (Variable): 真实标签。

        Returns:
            Variable: Softmax Cross Entropy 损失。
        """

        xp = cuda_module

        N = x.shape[0]
        log_z = utils.common.logsumexp(x, axis=1)
        log_p = x - log_z
        log_p = log_p[xp.arange(N), t.ravel()]
        y = -log_p.sum() / xp.float32(N)
        return y

    def backward(self, gy):
        """
        反向传播函数，计算相对于输入数据的梯度。

        Parameters:
            gy (Variable): 相对于前向传播结果的梯度。

        Returns:
            Variable: 相对于输入数据的梯度。
        """

        x, t = self.inputs
        N, CLS_NUM = x.shape

        gy.data = gy.data * 1 / N
        y = utils.functions_collect.softmax(x)
        # convert to one-hot
        t_onehot = cuda_module.eye(CLS_NUM, dtype=t.dtype)[t.data]
        y = (y - t_onehot) * gy
        return y

    def softmax_cross_entropy(self, x, t):
        """
        执行 Softmax Cross Entropy 操作。

        Parameters:
            x (Variable): 模型的输出结果。
            t (Variable): 真实标签。

        Returns:
            Variable: Softmax Cross Entropy 损失。
        """
        return self(x, t)

